tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py (1)

76-108: Deduplicate handle collection via get_weight_ipc_handles.

This keeps filtering/handle-building logic in one place and reduces drift between serialized and direct paths.

♻️ Suggested refactor

     def get_weight_ipc_handles_serialized(
         self,
         cuda_device: Optional[List[int]] = None,
         weight_filter: Optional[Callable[[str], bool]] = None,
     ):
         """
         Get base64-encoded serialized IPC handles for model weights.
@@
-        ret = {}
-        device_list = list(range(torch.cuda.device_count())) if cuda_device is None else cuda_device
-
-        for device in device_list:
-            all_handles = []
-            for item in self.all_weights[device]:
-                name, p = item
-                # Apply filter if provided
-                if weight_filter is not None and not weight_filter(name):
-                    continue
-                handle = reduce_tensor(p)
-                all_handles.append((name, handle))
-
-            # Serialize with base64-encoded pickle
-            serialized = base64.b64encode(pickle.dumps(all_handles)).decode("utf-8")
-            ret[self.device_uuid[device]] = serialized
+        ret = {}
+        handles = self.get_weight_ipc_handles(
+            cuda_device=cuda_device, weight_filter=weight_filter
+        )
+        for device_uuid, all_handles in handles.items():
+            # Serialize with base64-encoded pickle
+            serialized = base64.b64encode(pickle.dumps(all_handles)).decode("utf-8")
+            ret[device_uuid] = serialized
 
         return ret

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py`
around lines 76 - 108, The current get_weight_ipc_handles_serialized duplicates
the filtering and handle-building logic; replace the inner collection with a
call to the existing get_weight_ipc_handles so all filtering/handle creation is
centralized. Specifically, call
self.get_weight_ipc_handles(cuda_device=cuda_device,
weight_filter=weight_filter) to get the per-device mapping of (name, handle)
pairs (instead of iterating self.all_weights and calling reduce_tensor
directly), then for each device UUID serialize the returned list with pickle +
base64 and assign to ret[self.device_uuid[device]]; remove the duplicated loop
and direct reduce_tensor calls in get_weight_ipc_handles_serialized so it simply
wraps get_weight_ipc_handles results into base64-encoded pickles.

tensorrt_llm/serialization.py (1)

5-5: Avoid shared mutable defaults for approved_imports.

The regex-based module allowlist addition is a good step. One concern: approved_imports={} in Unpickler.__init__, load, and loads can leak mutations across calls. Prefer None defaults and initialize to {} inside.

Suggested update

 class Unpickler(pickle.Unpickler):

     def __init__(self,
                  *args,
-                 approved_imports={},
+                 approved_imports=None,
                  approved_module_patterns=None,
                  **kwargs):
         super().__init__(*args, **kwargs)
-        self.approved_imports = approved_imports
+        if approved_imports is None:
+            approved_imports = {}
+        self.approved_imports = approved_imports
         self.approved_module_patterns = approved_module_patterns or []
@@
 def load(file,
          *,
          fix_imports=True,
          encoding="ASCII",
          errors="strict",
          buffers=None,
-         approved_imports={},
+         approved_imports=None,
          approved_module_patterns=None):
+    if approved_imports is None:
+        approved_imports = {}
     return Unpickler(file,
                      fix_imports=fix_imports,
                      buffers=buffers,
                      encoding=encoding,
                      errors=errors,
                      approved_imports=approved_imports,
                      approved_module_patterns=approved_module_patterns).load()
@@
 def loads(s,
           /,
           *,
           fix_imports=True,
           encoding="ASCII",
           errors="strict",
           buffers=None,
-          approved_imports={},
+          approved_imports=None,
           approved_module_patterns=None):
+    if approved_imports is None:
+        approved_imports = {}
     if isinstance(s, str):
         raise TypeError("Can't load pickle from unicode string")

Also applies to: 130-154, 172-180, 190-201

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/serialization.py` at line 5, Change the mutable default for
approved_imports in Unpickler.__init__, Unpickler.load, and Unpickler.loads from
{} to None and initialize a fresh dict inside each method when None is passed
(e.g., if approved_imports is None: approved_imports = {}), so mutations don't
leak between calls; apply the same fix for any other functions in this module
that accept approved_imports (e.g., the other methods around the noted ranges)
to ensure each call gets its own dictionary.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro